# -*- coding: utf-8 -*-
"""
PBAR_ExaminePerformanceMergedMarkers.py

Created on Fri May 16 10:35:02 2014

@author: jkwong
"""

import glob, os, gc, objgraph, cPickle, time, copy, shutil
import numpy as np
import PBAR_Zspec, PBAR_Cargo
import matplotlib.pyplot as plt
from matplotlib import mpl
from scipy import ndimage
reload(PBAR_Zspec)

dataPath = r'E:\PBAR\data4\BasicScansStandardWidthPickle'
data0Path = r'E:\PBAR\data4\HighZFinder'
plotSaveDir = r'E:\PBAR\data4\HighZFinder'

num = 11

figureSize = (16, 10)

# Global parameters
MINBIN = 1
newWidth = 600
newHeight = 136
xCargoInternalBounds = np.array((38, 550))
yCargoInternalBounds = np.array((13, 106))
#cargoCountRange = np.array([0, 4e7]) # only consider parts of the zspec image that have rad below this value
cargoCountRange = np.array([0, 3e7]) # only consider parts of the zspec image that have rad below this value

ZspecCountRange = np.array([0, 275])

# make different size windows
# x and y
windowList = []
windowList.append(np.array((3, 1)))
windowList.append(np.array((5, 1)))
windowList.append(np.array((7, 1)))
windowList.append(np.array((4, 2)))
windowList.append(np.array((7, 2)))
# thresholds
discrimThresholdList = np.array((0, 2))
# fractions
fractionPixelsList = np.array((0.8, 1.0))

tolerance = [3, 1]

# get linear discriminant coefficients, polynomial fits

datasetDescription = PBAR_Zspec.ReadCargoDataDescriptionFile(r'E:\PBAR\data4\CargoSet2.txt')

numberCombinations =len(windowList) * len(discrimThresholdList) * len(fractionPixelsList) * len(datasetDescription['scanID'])

performanceSummary = {'datasetIndex':np.zeros(numberCombinations), 'discrimThreshold':np.zeros(numberCombinations), \
    'window': np.zeros((numberCombinations, 2)), 'fractionPixels': np.zeros(numberCombinations), \
    'f1': np.zeros(numberCombinations), 'success': np.zeros(numberCombinations), \
    'numMM': np.zeros(numberCombinations), 'numTruthMarkers': np.zeros(numberCombinations), \
    'numTrueHighZ': np.zeros(numberCombinations), 'TP': np.zeros(numberCombinations), \
    'numMMCorrect': np.zeros(numberCombinations), 'recall': np.zeros(numberCombinations), \
    'precision': np.zeros(numberCombinations), 'scanID': np.zeros(numberCombinations), \
    'config': np.zeros(numberCombinations)}

combinationIndex = -1
for datasetIndex, datasetDescript in enumerate(datasetDescription['scanID']):
    print('%d, %s' %(datasetIndex, datasetDescript))
    zspecScanNumber = np.int(datasetDescript)

    # Read Zspec
    filenameZspec = '%s-FDFC-All_SW.dat' %datasetDescription['scanID'][datasetIndex]
    fullFilenameZspec = os.path.join(dataPath, filenameZspec)
    
    # Read Cargo
    filenameCargo = 'PBAR-%s.cargoimageSW.dat' %datasetDescription['dataFile'][datasetIndex]
    fullFilenameCargo = os.path.join(dataPath, filenameCargo)

#    # read in the cargo image
#    print('Loading %s' %fullFilenameCargo)
#    with open(fullFilenameCargo, 'rb') as fid:
#        datCargoStandardWidth = cPickle.load(fid)    
#    datCargoSmall = PBAR_Zspec.CargoReduceSize(datCargoStandardWidth, newHeight, newWidth)
    
#    energy, datZspecStandardWidth = PBAR_Zspec.ReadZspecBasicScanPickle(fullFilenameZspec)
#    datZspecSmall = PBAR_Zspec.ZspecBasicReduceSize(datZspecStandardWidth, energy, newWidth)
#    # calculate the features for reduced size zspec image
#    discrimSmall = PBAR_Zspec.CalculateFeaturesZspecBasic(datZspecSmall, energy, MINBIN)    
    
    # Read in marker files
    filenameMarker = filenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')
    fullFilenameMarker = fullFilenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')
    
    # some don't have marker files
    if os.path.exists(fullFilenameMarker):
        markerStandardWidth0 = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
    else:
        markerStandardWidth0 = []
    
#     Modify the marker files for truncated
#    multiplier = newWidth/float(datZspecStandardWidth.shape[0])
    multiplier = newWidth/float(newWidth)

    offset = 0.0
    markerStandardWidth = PBAR_Zspec.ModifyMarkersXPosition(markerStandardWidth0, multiplier, offset)

    numberMarkers = len(markerStandardWidth)
    
    # determine if marker is high z (DU or W) or not (Fe or Pb)
    highZMaterial = np.zeros(numberMarkers) == 1 # start with all false
    for markerIndex, marker in enumerate(markerStandardWidth):
        # s = special, w = tungsten, du = depleted uranium
        if ( (marker['target'][0].lower() == 's') or (marker['target'][0].lower() == 'w') or (marker['target'][0].lower() == 'd')):
            highZMaterial[markerIndex] = True

    for windowIndex, window in enumerate(windowList):
        for discrimThresholdIndex, discrimThreshold in enumerate(discrimThresholdList):
            for fractionPixelsIndex, fractionPixels in enumerate(fractionPixelsList):
                
                combinationIndex += 1

                tStart = time.time()
                
                # makes them all have the same reference, I think
#                highZMarkersCrossRef = [{'mergedMergedIndex':[]}] * numberMarkers

                highZMarkersCrossRef = []
                for i in xrange(numberMarkers):
                    highZMarkersCrossRef.append({})
                    highZMarkersCrossRef[i]['mergedMergedIndex'] = []
                
                fullFilenamePotential = \
                    fullFilenameMarker.replace('cargomarkerSW', \
                        '%d_%d_%d_%3.2f.potential' %(window[0], window[1], discrimThreshold, fractionPixels) )
                print('Working on %s' %fullFilenamePotential)
                
                # load original high markers
                with open(fullFilenamePotential, 'rb') as fid:
                    print('Load Original Markers: %s' %fullFilenamePotential)
                    highZMarkersList = cPickle.load(fid)
                # load the merged markers
                fullFilenameMarkerMerged = fullFilenamePotential.replace('potential', 'highZMerged')
                with open(fullFilenameMarkerMerged, 'rb') as fid:
                    print('Load Merged Markers: %s' %fullFilenameMarkerMerged)
                    highZMergedMarkersList = cPickle.load(fid)
                # read merged merged markers
                fullFilenameMarkerMerged = fullFilenamePotential.replace('potential', 'highZMergedMerged')
                with open(fullFilenameMarkerMerged, 'rb') as fid:
                    print('Load Merged Merged Markers: %s' %fullFilenameMarkerMerged)
                    highZMergedMergedMarkersList = cPickle.load(fid)
                # read the merged merged marker performance
                fullFilenamePerformance = fullFilenamePotential.replace('potential', 'performance')
                with open(fullFilenamePerformance, 'rb') as fid:
                    print('Write: %s' %fullFilenamePerformance)
                    dat = cPickle.load(fid)
                print('time elapsed3: %3.3f' %(time.time() - tStart))

#                break
#            break
#        break
#    break

# Load from file
fullFilenamePerformance = os.path.join(data0Path, 'performanceSummary.dat')
with open(fullFilenamePerformance, 'rb') as fid:
    print('Load: %s' %fullFilenamePerformance)
    performanceSummary = cPickle.load(fid)


#  Print average metrics for each window, threshold and fractionpixels

metricName = 'f1'
metricNameList = ['success', 'f1', 'precision', 'recall', 'TP']

for metricName in metricNameList:
    for discrimThresholdIndex, discrimThreshold in enumerate(discrimThresholdList):
        cutt = (performanceSummary['discrimThreshold']==discrimThreshold) & ~np.isnan(performanceSummary[metricName])
        print('mean %s[discrimThresold == %3.2f]: %3.3f'  %(metricName, discrimThreshold, performanceSummary[metricName][cutt].mean()) )
    
    for windowIndex, window in enumerate(windowList):
        cutt = (performanceSummary['window'][:,0]==window[0]) & (performanceSummary['window'][:,1]==window[1]) & ~np.isnan(performanceSummary[metricName])
        print('mean %s[window == [%d,%d]]: %3.3f'  %(metricName, window[0], window[1], performanceSummary[metricName][cutt].mean()) )
    
    for fractionPixelsIndex, fractionPixels in enumerate(fractionPixelsList):
        cutt = (performanceSummary['fractionPixels']==fractionPixels) & ~np.isnan(performanceSummary[metricName])
        print('mean %s[fractionPixels == %3.2f]: %3.3f'  %(metricName, fractionPixels, performanceSummary[metricName][cutt].mean()) )

# for the best combination for each metric
metricNameList = ['success', 'f1', 'precision', 'recall', 'TP']
metricValuesMean = {}
metricValuesMeanList = {}
for metricName in metricNameList:

    metricValuesMean[metricName] = np.zeros((len(discrimThresholdList), len(windowList) , len(fractionPixelsList) ))
    metricValuesMeanList[metricName] = np.zeros(len(discrimThresholdList)*len(windowList)*len(fractionPixelsList))

    # best settings
    bestValue = -1
    discrimThresholdIndexBest = 0
    windowIndexBest = 0
    fractionPixelsIndexBest = 0
    n = -1

    for discrimThresholdIndex, discrimThreshold in enumerate(discrimThresholdList):
        for windowIndex, window in enumerate(windowList):
            for fractionPixelsIndex, fractionPixels in enumerate(fractionPixelsList):
                n += 1
                cutt = (performanceSummary['discrimThreshold']==discrimThreshold) & \
                    (performanceSummary['window'][:,0]==window[0]) & \
                    (performanceSummary['fractionPixels']==fractionPixels) & \
                    ~np.isnan(performanceSummary[metricName])
                metricValue = performanceSummary[metricName][cutt].mean()

                metricValuesMean[metricName][discrimThresholdIndex,windowIndex,fractionPixelsIndex ] = metricValue
                metricValuesMeanList[metricName][n] = metricValue
                if metricValue > bestValue:
                    discrimThresholdIndexBest = copy.copy(discrimThresholdIndex)
                    windowIndexBest = copy.copy(windowIndex)
                    fractionPixelsIndexBest = copy.copy(fractionPixelsIndex)
                    bestValue = copy.copy(metricValue)
    discrimThreshold = discrimThresholdList[discrimThresholdIndexBest]
    window = windowList[windowIndexBest]
    fractionPixels = fractionPixelsList[fractionPixelsIndexBest]    
    print('best %s[%3.2f, [%d, %d], %3.2f ]: %3.4f'  %(metricName, discrimThreshold, window[0], window[1], fractionPixels, bestValue) )

# Order combinations by each metric
# - basically more general than the previous section of code which only displays the top rank

discrimThresholdListList = np.zeros(len(discrimThresholdList)* len(windowList) * len(fractionPixelsList))
windowListList =np.zeros((len(discrimThresholdList)* len(windowList) * len(fractionPixelsList), 2))
fractionPixelsListList = np.zeros(len(discrimThresholdList)* len(windowList) * len(fractionPixelsList))
n = -1
for discrimThresholdIndex, discrimThreshold in enumerate(discrimThresholdList):
    for windowIndex, window in enumerate(windowList):
        for fractionPixelsIndex, fractionPixels in enumerate(fractionPixelsList):
            n +=1
            discrimThresholdListList[n] = discrimThreshold
            windowListList[n,:] = window
            fractionPixelsListList[n] = fractionPixels

# present the top settings for each
topN = 5
for metricName in metricNameList:
    sortedOrder = np.argsort(metricValuesMeanList[metricName])[-1::-1]
    for i in xrange(topN):
        n = sortedOrder[i]
        discrimThreshold = discrimThresholdListList[n]
        window = windowListList[n]
        fractionPixels = fractionPixelsListList[n]
        bestValue = metricValuesMeanList[metricName][n]
        print('%d best %s[%3.2f, [%d, %d], %3.2f ]: %3.4f'  %(i, metricName, discrimThreshold, window[0], window[1], fractionPixels, bestValue) )

# present the top settings for each
topN = 5
for metricName in metricNameList:
    sortedOrder = np.argsort(metricValuesMeanList[metricName])
    for i in xrange(topN):
        n = sortedOrder[i]
        discrimThreshold = discrimThresholdListList[n]
        window = windowListList[n]
        fractionPixels = fractionPixelsListList[n]
        bestValue = metricValuesMeanList[metricName][n]
        print('%d worst %s[%3.2f, [%d, %d], %3.2f ]: %3.4f'  %(i, metricName, discrimThreshold, window[0], window[1], fractionPixels, bestValue) )


# find the highest values


performanceSummary['numTruthMarkers'] > 5

discrimThreshold == discrimThresholdList[0]
window = windowList[0]
fractionPixels == fractionPixelsList[0]

metricName = metricNameList[1]
cutt = (performanceSummary['discrimThreshold']==discrimThreshold) & \
    (performanceSummary['window'][:,0]==window[0]) & \
    (performanceSummary['window'][:,1]==window[1]) & \
    (performanceSummary['fractionPixels']==fractionPixels) & \
    ~np.isnan(performanceSummary[metricName])

metricValue = performanceSummary[metricName][cutt]
sortOrder = np.argsort(metricValue)[-1::-1]

datasetIndices = np.where(cutt)[0][sortOrder]

for i in datasetIndices:
    index = performanceSummary['datasetIndex'][i]
    window = performanceSummary['window'][i]
    fractionPixels = performanceSummary['fractionPixels'][i]
    discrimThreshold = performanceSummary['discrimThreshold'][i]
        
    filename =  '%d_%d_%d_%3.2f.potential' %(window[0], window[1], discrimThreshold, fractionPixels) 
    print('%s, %3.3f' %(filename, performanceSummary[metricName][i]))

# copy some of the files to a temporary location for faster viewing

discrimThreshold == discrimThresholdList[0]
window = windowList[0]
fractionPixels == fractionPixelsList[0]

metricName = metricNameList[1]
cutt = (performanceSummary['discrimThreshold']==discrimThreshold) & \
    (performanceSummary['window'][:,0]==window[0]) & \
    (performanceSummary['window'][:,1]==window[1]) & \
    (performanceSummary['fractionPixels']==fractionPixels) & \
    ~np.isnan(performanceSummary[metricName])

metricValue = performanceSummary[metricName][cutt]
sortOrder = np.argsort(metricValue)[-1::-1]


sourcePlotDir = r'E:\PBAR\data4\HighZFinder'
tempPlotDir = r'E:\PBAR\data4\TempPlots'

for i in datasetIndices:
    index = np.int(performanceSummary['datasetIndex'][i])  
    scanIdStr = datasetDescription['scanID'][index]
    window = performanceSummary['window'][i]
    fractionPixels = performanceSummary['fractionPixels'][i]
    discrimThreshold = performanceSummary['discrimThreshold'][i]

#    6960_7_2_2_1.00_Cargo_MM.png
    filename = '%s_%d_%d_%d_%3.2f_Cargo_MM.png' %(scanIdStr, window[0], window[1], discrimThreshold, fractionPixels)
    print('Copying: %s' %filename)
    fullFilenameSource = os.path.join(sourcePlotDir, filename)
    fullFilenameDestination = os.path.join(tempPlotDir, filename)
    shutil.copy(fullFilenameSource, fullFilenameDestination)



# find image of best performance

plt.figure()

plt.plot(performanceSummary['discrimThreshold'], performanceSummary['f1'], '.k')
plt.plot(performanceSummary['discrimThreshold'], performanceSummary['f1'], '.k')

